import cv2
import os

# 路径英文
train_dir = r"D:\zjk\result_texture"
gt_dir = r"D:\zjk\res"
save_dir = "./my_data"


def gen_train_image():
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    if not os.path.exists(os.path.join(save_dir, "train")):
        os.mkdir(os.path.join(save_dir, "train"))
    if not os.path.exists(os.path.join(save_dir, "train_mask")):
        os.mkdir(os.path.join(save_dir, "train_mask"))
    for file in os.listdir(train_dir):
        gt_name = file.split('_')[0]
        gt_path = os.path.join(gt_dir, gt_name + ".bmp")
        train_path = os.path.join(train_dir, file)

        img = cv2.imread(train_path)
        cv2.imwrite(os.path.join(save_dir, "train", file.replace('bmp', 'jpg')),img)
        img = cv2.imread(gt_path)
        cv2.imwrite(os.path.join(save_dir, "train_mask", file.replace('bmp', 'png')), img)


if __name__ == '__main__':
    gen_train_image()